package com.lee.controller;

import cn.hutool.core.util.IdUtil;
import com.lee.service.SSEService;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.schedulers.Schedulers;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;

@RestController
@RequestMapping("test")
public class TestController {
    @Autowired
    private SSEService sseService;
    @Autowired
    private ArkService service;
    @GetMapping("createSse")
    public SseEmitter createSse(String uid)
    {
        return sseService.crateSse(uid);
    }

    @GetMapping("sendMsg")
    public String sseChat(String uid) throws InterruptedException {
        for (int i = 0; i < 10; i++) {
            sseService.sendMessage(uid,"消息"+i,IdUtil.fastUUID().replace("-",""));
            Thread.sleep(1000);
        }
        return "OK";
    }
    @GetMapping("closeSse")
    public void closeSse(String uid){
        sseService.closeSse(uid);
    }

    @GetMapping("abstract")
    public String getAbstract() throws InterruptedException {

        System.out.println("\n----- streaming request -----");
        final List<ChatMessage> streamMessages = new ArrayList<>();
        final ChatMessage streamSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM)
                .content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build();
        final ChatMessage streamUserMessage = ChatMessage.builder().role(ChatMessageRole.USER)
                .content("常见的十字花科植物有哪些？").build();
        streamMessages.add(streamSystemMessage);
        streamMessages.add(streamUserMessage);
        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20250207210413-l6c9q")
                .messages(streamMessages)
                .build();

        service.streamChatCompletion(streamChatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .subscribeOn(Schedulers.io())
                .observeOn(Schedulers.single())
                .subscribe(
                        choice -> {
                            if (choice.getChoices().size() > 0) {
                                sseService.sendMessage("1",
                                        "消息"+ IdUtil.fastUUID().replace("-",""),
                                        choice.getChoices().get(0).getMessage().getContent()+"");
                            }
                        },
                        Throwable::printStackTrace,
                        () -> {
                            // 流式输出完成后的处理
                            System.out.println("=========== 流式输出完成! ===========");
                        }
                );

        // just wait for result
        // shutdown service
//        service.shutdownExecutor();
        return "ok";
    }
}